{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce655333",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "from ultralytics import YOLO\n",
    "\n",
    "from ppocronnx import TextSystem\n",
    "import json\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff2f111f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def im_show(image):\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    plt.imshow(image)\n",
    "    plt.show()\n",
    "    \n",
    "CLASSES = ['number']\n",
    "\n",
    "COLORS = [(255, 0, 0)]\n",
    "\n",
    "def draw_boxes(image, boxes, rec_results):\n",
    "    mask_img = image.copy()\n",
    "    font = cv2.FONT_HERSHEY_COMPLEX\n",
    "    for box,number_score in zip(boxes, rec_results):\n",
    "        color = COLORS[0]\n",
    "        x1, y1, x2, y2 = box\n",
    "        cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, 5)\n",
    "        number, score = number_score\n",
    "        cv2.putText(mask_img,str(number),(x1-10,y1-10),font,3,color,5)\n",
    "    return mask_img\n",
    "\n",
    "def detect_image(detector, image):\n",
    "    #print(\"=== come into detect frame \")\n",
    "    results = detector(image)\n",
    "    #print(\"=== image detected\")\n",
    "    boxes = []\n",
    "    crop_number_imgs = []\n",
    "    for box in results[0].boxes:\n",
    "        xyxy = box.xyxy[0].cpu().numpy().tolist()\n",
    "        x1, y1, x2, y2 = [int(_) for _  in xyxy]\n",
    "        boxes.append([x1, y1, x2, y2])\n",
    "        crop_img = image[y1:y2, x1:x2, :]\n",
    "        crop_number_imgs.append(crop_img)\n",
    "    return boxes, crop_number_imgs\n",
    "\n",
    "def detect_frame(detector, frame, vertical_axes):\n",
    "    v0, v1 = vertical_axes\n",
    "    frame = frame[v0:v1, :, :]\n",
    "    \n",
    "    h,w, _ = frame.shape\n",
    "    print(frame.shape)\n",
    "    frame_boxes = []\n",
    "    frame_crop_imgs = []\n",
    "    for vsplit in range(0, h, 640):\n",
    "        for hsplit in range(0, w, 640):\n",
    "            split_img = frame[vsplit:vsplit+640, hsplit:hsplit+640, :]\n",
    "            print(\"vsplit:\", vsplit, \"hsplit:\",hsplit, split_img.shape)\n",
    "            boxes, crop_num_imgs = detect_image(detector, split_img)\n",
    "            frame_crop_imgs.extend(crop_num_imgs)\n",
    "            img_boxes = []\n",
    "            for box in boxes:\n",
    "                x1, y1, x2, y2 = box\n",
    "                img_box = [x1+hsplit, y1+vsplit+v0, x2+hsplit, y2+vsplit+v0]\n",
    "                img_boxes.append(img_box)\n",
    "            frame_boxes.extend(img_boxes)\n",
    "    return frame_boxes,frame_crop_imgs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ed9307a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = 'models/yolov11_number_best_v1.onnx'\n",
    "number_detector = YOLO(model_path, task='detect')\n",
    "\n",
    "rec_model = 'models/number_ocr_model/number_rec.onnx'\n",
    "text_sys = TextSystem(rec_model_path=rec_model, ort_providers=['CUDAExecutionProvider'])\n",
    "\n",
    "\n",
    "frame = cv2.imread('cba_frame_0.jpg')\n",
    "\n",
    "vertical_axes = [750, 1390]\n",
    "boxes, crop_number_imgs = detect_frame(number_detector, frame, vertical_axes)\n",
    "#print(\"== detect ok \")\n",
    "number_results = text_sys.ocr_lines(crop_number_imgs)\n",
    "#print(\"== rec finish!\")\n",
    "out_img = draw_boxes(frame, boxes, number_results)\n",
    "\n",
    "im_show(out_img)\n",
    "\n",
    "cv2.imwrite(\"number_frame.jpg\", out_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81ed1b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = 'models/yolov11_number_best_v1.onnx'\n",
    "number_detector = YOLO(model_path, task='detect')\n",
    "\n",
    "rec_model = 'models/number_ocr_model/number_rec.onnx'\n",
    "text_sys = TextSystem(rec_model_path=rec_model, ort_providers=['CUDAExecutionProvider'])\n",
    "#text_sys = TextSystem(rec_model_path=rec_model, ort_providers=['CPUExecutionProvider'])\n",
    "vertical_axes = [750, 1390]\n",
    "\n",
    "v_idx = 0\n",
    "#input_video = '/home/imvision/hy/samurai-master/data/output%03d.mp4' % v_idx\n",
    "input_video = '/home/imvision/hy/samurai-master/data/output_0_to_20.mp4'\n",
    "#input_video = './datas/output000.mp4'\n",
    "reader = cv2.VideoCapture(input_video)\n",
    "\n",
    "width = int(reader.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height = int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "\n",
    "#output_video = 'datas/v%03d_number_out.mp4' % v_idx\n",
    "output_video = 'datas/v0_to_20_number_out2.mp4'\n",
    "#output_video = 'datas/v0_number_out2.mp4'\n",
    "writer = cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*\"mp4v\"), 25, (width, height))\n",
    "more = True\n",
    "frame_idx = -1\n",
    "\n",
    "frame_2_number_box = {}\n",
    "#output_boxes_path = 'datas/v%03d_number_out.json' % v_idx\n",
    "output_boxes_path = 'datas/v0_to_20_number_out2.json'\n",
    "#output_boxes_path = 'datas/v0_number_out2.json'\n",
    "\n",
    "def write_boxes(boxes_path, frame_2_boxes):\n",
    "    with open(boxes_path, 'w') as fout:\n",
    "        for frame_idx, frame_boxes in frame_2_boxes.items():\n",
    "            box_str = json.dumps(frame_boxes)\n",
    "            fout.write(\"{}:{}\\n\".format(frame_idx, box_str))\n",
    "\n",
    "tic = time.time()\n",
    "while more:\n",
    "    more, frame = reader.read()\n",
    "    frame_idx += 1\n",
    "    if not more:\n",
    "        break\n",
    "#     if frame_idx >=10:\n",
    "#         break\n",
    "    if frame_idx % 5 == 0:\n",
    "        toc = time.time()\n",
    "        print(\"== handle \", frame_idx, \" fps \", 20/(toc-tic))\n",
    "        tic = toc\n",
    "    #print(\"== before detect\")\n",
    "    boxes, crop_number_imgs = detect_frame(number_detector, frame, vertical_axes)\n",
    "    #print(\"== detect ok \")\n",
    "    number_results = text_sys.ocr_lines(crop_number_imgs)\n",
    "    #print(\"== rec finish!\")\n",
    "    out_img = draw_boxes(frame, boxes, number_results)\n",
    "    frame_boxes = []\n",
    "    for box, number_score in zip(boxes, number_results):\n",
    "        number, score = number_score\n",
    "        frame_boxes.append((box, number, float(score)))\n",
    "    frame_2_number_box[frame_idx] = frame_boxes\n",
    "    #im_show(out_img)\n",
    "    writer.write(out_img)\n",
    "\n",
    "write_boxes(output_boxes_path, frame_2_number_box)\n",
    "reader.release()\n",
    "writer.release()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
